import numpy as np
from scipy.integrate import solve_ivp
from math import sin, cos, radians, atan, degrees
import matplotlib.pyplot as plt

# --- Matplotlib 中文设置 ---
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Heiti TC'] 
plt.rcParams['axes.unicode_minus'] = False 

# --- I. 物理常数和弹丸参数 ---
RHO0 = 1.225; T0 = 288.15; L = 0.0065
G = 9.80; R = 287.05
P_TERM = (G / (R * L)) - 1 
C_SOUND = 338.0
M = 5.0; A = 0.0095

# --- II. 目标和搜索参数 ---
TARGET_RANGE = 1200.0  # 目标射程 (m)
TARGET_FALL_ANGLES = [7.0, 45.0, 60.0] # 目标落地角 (度)
RANGE_TOL = 0.5        # 射程误差容忍度
VEL_TOL = 0.1          # 初速间隔容忍度
ANGLE_TOL = 0.01       # 射角搜索容忍度 (度)

# --- III. Cd-Ma 和 Rho-z 关系函数 (不变) ---

MA_POINTS = np.array([0.0, 0.7, 1.0, 1.3, 2.0, 2.5, 10.0]) 
CD_POINTS = np.array([0.48, 0.45, 0.95, 0.45, 0.40, 0.40, 0.40])

def get_drag_coefficient(v, c_sound):
    mach_number = v / c_sound
    Cd = np.interp(mach_number, MA_POINTS, CD_POINTS)
    return Cd

def get_air_density(z):
    h = max(0.0, z)
    temp_ratio = (T0 - L * h) / T0
    if temp_ratio <= 0:
        return 0.0
    rho = RHO0 * (temp_ratio) ** P_TERM
    return rho

# --- IV. 弹道微分方程 (不变) ---

def equations_of_motion(t, S, theta, v0, params):
    x, z, vx, vz = S
    m, A, g, c_sound = params
    v = np.sqrt(vx**2 + vz**2)
    rho = get_air_density(z)
    Cd = get_drag_coefficient(v, c_sound)
    drag_coeff = -(Cd * rho * A) / (2.0 * m) 
    ax = drag_coeff * v * vx
    az = drag_coeff * v * vz - g
    return [vx, vz, ax, az]

# --- V. 轨迹积分函数 (修改：返回落地角和速度分量) ---

def calculate_trajectory(v0, theta):
    """
    积分弹道，返回 R_final, fall_angle, vx_final, vz_final, 轨迹点
    """
    params = (M, A, G, C_SOUND)
    theta_rad = radians(theta)
    S0 = [0.0, 0.0, v0 * cos(theta_rad), v0 * sin(theta_rad)]
    
    def hit_ground(t, S, theta, v0, params):
        return S[1]
    
    hit_ground.terminal = True 
    hit_ground.direction = -1 
    t_span = [0, 500] 
    
    sol = solve_ivp(equations_of_motion, t_span, S0, 
                    args=(theta, v0, params), method='RK45', 
                    events=hit_ground, rtol=1e-7, atol=1e-9, dense_output=True)
    
    if len(sol.y_events[0]) > 0:
        # 落地状态
        final_state = sol.y_events[0][0]
        final_x = final_state[0]
        vx_final = final_state[2]
        vz_final = final_state[3]
        
        # 计算落地角 alpha (与x轴正向的夹角，取绝对值)
        fall_angle_rad = atan(abs(vz_final / vx_final))
        fall_angle_deg = degrees(fall_angle_rad)
        
        # 提取轨迹点
        t_hit = sol.t_events[0][0]
        t_plot = np.linspace(0, t_hit, 300)
        z_plot = sol.sol(t_plot)
        
        return final_x, fall_angle_deg, vx_final, vz_final, z_plot[0], z_plot[1]
    else:
        return -1.0, 999.0, 0.0, 0.0, None, None

# --- VI. 核心双重搜索函数 ---

def find_required_velocity(theta):
    """
    内层搜索：给定射出角 theta，使用二分法寻找满足 TARGET_RANGE 的 v0。
    返回: (v0_found, R_found, Fall_Angle)
    """
    v_min = 10.0  # 搜索下限
    v_max = 500.0  # 搜索上限
    
    # 确保搜索区间包含 TARGET_RANGE
    R_max, _, _, _, _, _ = calculate_trajectory(v_max, theta)
    R_min, _, _, _, _, _ = calculate_trajectory(v_min, theta)
    
    # 尝试自动调整 v_min
    if R_min > TARGET_RANGE:
        v_min_auto = v_min
        while R_min > TARGET_RANGE and v_min_auto > 1.0:
            v_min_auto /= 1.5
            R_min, _, _, _, _, _ = calculate_trajectory(v_min_auto, theta)
        v_min = v_min_auto
    
    if R_min > TARGET_RANGE or R_max < TARGET_RANGE:
        # print(f"警告：θ={theta:.2f}°，目标射程不在 [{v_min:.1f}, {v_max:.1f}] 范围内。")
        return -1.0, -1.0, 999.0

    while (v_max - v_min > VEL_TOL):
        v_mid = (v_min + v_max) / 2.0
        R_mid, _, _, _, _, _ = calculate_trajectory(v_mid, theta)
        
        if R_mid < 0: # 积分失败
             v_max = v_mid
             continue

        if abs(R_mid - TARGET_RANGE) < RANGE_TOL:
            break
            
        if R_mid > TARGET_RANGE:
            v_max = v_mid
        else:
            v_min = v_mid
            
    v_found = (v_min + v_max) / 2.0
    R_found, fall_angle, _, _, _, _ = calculate_trajectory(v_found, theta)

    return v_found, R_found, fall_angle


def find_launch_angle_and_velocity(target_fall_angle):
    """
    外层搜索：寻找满足目标落地角的射出角 theta。
    """
    print(f"\n--- 🔎 搜索目标落地角 α = {target_fall_angle:.1f}° ---")
    
    # 搜索范围 (从平射到高射)
    theta_min = 1.0
    theta_max = 89.0
    
    # 迭代搜索 theta 的过程 (使用二分法近似)
    i = 0
    while (theta_max - theta_min > ANGLE_TOL) and i < 50:
        i += 1
        theta_mid = (theta_min + theta_max) / 2.0
        
        # 内层搜索：找到当前 theta 下满足射程的 v0 和对应的落地角 alpha
        v0, R, alpha = find_required_velocity(theta_mid)
        
        if v0 < 0: # v0 搜索失败
            # 假设失败是由于 theta 范围不合理，缩小范围
            if theta_mid < 45: theta_min = theta_mid
            else: theta_max = theta_mid
            continue
            
        alpha_error = abs(alpha - target_fall_angle)
        
        print(f"迭代 {i:2}: θ={theta_mid:.3f}°, v₀={v0:.3f} m/s, α={alpha:.3f}°, 误差={alpha_error:.3f}°")

        if alpha_error < ANGLE_TOL:
            break
            
        # 二分法更新 theta 搜索范围
        # 通常，射出角 theta 越高，落地角 alpha 越接近 90度（更陡峭）
        if alpha > target_fall_angle:
            # 落地角太陡峭，需要降低射出角 theta
            theta_max = theta_mid
        else:
            # 落地角太平缓，需要增加射出角 theta
            theta_min = theta_mid

    # 最终结果
    theta_found = (theta_min + theta_max) / 2.0
    v0_found, R_found, alpha_found = find_required_velocity(theta_found)
    
    print(f"✅ 找到解: θ={theta_found:.3f}°, v₀={v0_found:.3f} m/s, α={alpha_found:.3f}°")

    if abs(alpha_found - target_fall_angle) > ANGLE_TOL * 2:
        print("警告：最终落地角误差较大，可能需要更多迭代或更小的容忍度。")
        
    return theta_found, v0_found, R_found, alpha_found

# --- VII. 执行计算、绘制和总结 ---

results = []
for target_angle in TARGET_FALL_ANGLES:
    theta, v0, R, alpha = find_launch_angle_and_velocity(target_angle)
    if v0 > 0:
        results.append({'alpha': alpha, 'theta': theta, 'v0': v0})

if not results:
    print("\n❌ 未能找到任何满足条件的解，请检查搜索范围和容忍度。")
else:
    # 绘制所有轨迹
    plt.figure(figsize=(12, 7))
    colors = ['r', 'b', 'g']
    
    print("\n--- 📈 绘制轨迹图 ---")
    
    for i, res in enumerate(results):
        # 积分得到轨迹数据
        _, _, _, _, x_traj, z_traj = calculate_trajectory(res['v0'], res['theta'])
        
        if x_traj is not None:
            label = f'$\\alpha={res["alpha"]:.1f}^\circ$, $\\theta={res["theta"]:.1f}^\circ$, $v_0={res["v0"]:.1f} \mathrm{{m/s}}$'
            plt.plot(x_traj, z_traj, label=label, color=colors[i], linewidth=2)
            
            # 标记最大射高
            max_height = np.max(z_traj)
            max_height_idx = np.argmax(z_traj)
            plt.plot(x_traj[max_height_idx], max_height, colors[i] + 'x', markersize=8, mew=2)
            
            print(f"轨迹 {i+1}: α={res['alpha']:.1f}°，v₀={res['v0']:.3f} m/s，Z_max={max_height:.3f} m")

    # 绘制目标点和地面线
    plt.plot(TARGET_RANGE, 0, 'ko', markersize=8, label=f'Target ({TARGET_RANGE:.0f} m)')
    plt.hlines(0, 0, TARGET_RANGE * 1.1, color='gray', linestyle='--', linewidth=0.8, label='Ground')
    
    plt.title(r'Trajectory that meets the target range of $R=1200 \,\mathrm{m}$', fontsize=16)
    plt.xlabel('Horizontal distance x (m)', fontsize=14)
    plt.ylabel('Height z (m)', fontsize=14)
    plt.grid(True, linestyle=':', alpha=0.6)
    plt.legend(title='Target impact angle, Launch angle, Initial velocity')
    plt.gca().set_aspect('equal', adjustable='box')
    plt.ylim(bottom=0)
    plt.xlim(right=TARGET_RANGE * 1.1)
    
    plt.show()

# 总结结果表格
print("\n--- 📊 最终结果总结 ---")
print("| 目标落地角 (α_target) | 射出角 (θ) | 初速度 (v₀, m/s) | 实际落地角 (α) |")
print("| :---: | :---: | :---: | :---: |")
for res in results:
    print(f"| {res['alpha']:.1f}° | {res['theta']:.3f}° | {res['v0']:.3f} | {res['alpha']:.3f}° |")